# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import datetime
import functools
import logging
import signal
from collections import namedtuple

from absl import app, flags
from tornado.ioloop import IOLoop

import coin.flow.subscriber as flow_subscriber
from coin.exchange.base.order_gateway import OrderSide, OrderType
from coin.exchange.binance.kr_rest.currency import BinanceCurrency
from coin.exchange.binance.kr_rest.product import BinanceProduct
from coin.exchange.kr_rest.product_holders import get_holder_from_product
from coin.exchange.shared.kr_rest.currency import SharedCurrency
from coin.exchange.shared.kr_rest.product import SharedProduct
from coin.flow.simple_queue import SimpleQueue
from coin.strategy.executor.util.cancel_all_executor import (SimpleCancelAllExecutor)
from coin.strategy.hk_hard_arb.balance_printer1 import (ArbInfoPnlPrinter, PassAggArbProducts)
from coin.strategy.hk_hard_arb.price.pass_agg import PriceCalculator
from coin.strategy.hk_hard_arb.sizer.pass_agg import SingleProductSizeCalculator
from coin.strategy.hk_hard_arb.util import (PassAgg, PassAggExchangeMap, load_clean_json)
from coin.strategy.mm.base.multi_feed_subsystem import MultiFeedSubsystem
from coin.strategy.mm.feed import SubscriptionRequest, get_flow_book_builder_cls
from coin.strategy.mm.live_envs import get_live_env_cls
from coin.strategy.executor.passive.factory import construct_mm_executor
from coin.strategy.executor.passive.util import (convert_executor_config_from_dict)

from coin.strategy.hk_hard_arb.og_util import print_working_orders

FLAGS = flags.FLAGS

flags.DEFINE_string('subreq_passive', 'Subscription request - passive', '')

flags.DEFINE_string('subreq_aggressive', 'Subscription request - aggressive', '')

flags.DEFINE_string('og_config_passive', '', '')

flags.DEFINE_string('og_config_aggressive', '', '')

flags.DEFINE_string('config_filename', '', '')

flags.DEFINE_integer('exit_after_min', 30, '')
flags.DEFINE_boolean('dry_run', False, '')

PendingAggression = namedtuple('PendingAggression', ['fill_price', 'fill_qty', 'timestamp'])


def agg_order_side(order):
  if order.order_side == OrderSide.SELL:
    order_side_agg = OrderSide.BUY
  elif order.order_side == OrderSide.BUY:
    order_side_agg = OrderSide.SELL
  else:
    raise ValueError('Invalid side: %s' % str(order.order_side))
  return order_side_agg


class SingleProductSubStrategy(object):
  def __init__(
      self,
      arb_product,
      parent,  # TODO(inkyu): Separate me
      config,
      logger=None):
    self._arb_product = arb_product
    self._product_holders = PassAggExchangeMap(
        passive=get_holder_from_product(arb_product.native_passive),
        aggressive=get_holder_from_product(arb_product.native_aggressive))

    self._logger = logger

    # Execution params
    self._edge = float(config['edge_bp']) / 10000.
    self._agg_edge = float(config['agg_edge_bp']) / 10000.
    self._pull_edge = None
    if 'pull_edge_bp' in config:
      self._pull_edge = float(config['pull_edge_bp']) / 10000.
    self._reserve_balance = config['max_pos'] + config['min_pos']
    self._lot_size = float(config['lot_size'])
    self._lot_size_min_agg = float(config['lot_size_min_agg'])
    self._lot_size_max_agg = float(config['lot_size_max_agg'])

    self._books = PassAggExchangeMap(None, None)

    # composite instances
    self._parent = parent

    book_askt_func = (
        lambda book: book.get_notional_asks_by_amt(config['price_det_notional']).entries[-1].price)
    book_bidt_func = (
        lambda book: book.get_notional_bids_by_amt(config['price_det_notional']).entries[-1].price)
    self._pricer = PriceCalculator(basis_ma_window=(int(config['basis_ma_window_min'] * 60e9)),
                                   product_holders=self._product_holders,
                                   book_askt_funcs=PassAggExchangeMap(book_askt_func,
                                                                      book_askt_func),
                                   book_bidt_funcs=PassAggExchangeMap(book_bidt_func,
                                                                      book_bidt_func))

    delay_aggressive_order = 5e9
    if 'delay_aggressive_order_sec' in config:
      delay_aggressive_order = int(config['delay_aggressive_order_sec'] * 1e9)

    self._sizer = SingleProductSizeCalculator(base_reserve=self._reserve_balance,
                                              lot_size=self._lot_size,
                                              delay_aggressive_order=delay_aggressive_order)

    self._parent = parent
    self._executors = PassAggExchangeMap(None, None)
    self._config = config

    # Aggression execution params
    self._last_timestamp_agg_buy = 0
    self._last_timestamp_agg_sell = 0

    self._pending_aggressions = {OrderSide.BUY: None, OrderSide.SELL: None}

    # TODO(inkyu): Wrap this
    self._parent._balance_printer._reserves[self._arb_product.product.base] = (
        self._reserve_balance)

  @property
  def _product_passive(self):
    return self._arb_product.native_passive

  @property
  def _product_aggressive(self):
    return self._arb_product.native_aggressive

  def on_og_reset(self):
    # Intialize passive executor
    self._logger.info('Initializing passive executor...')

    # TODO(inkyu): Make proto & separate execution parameters
    execution_params_passive = convert_executor_config_from_dict(self._product_passive.symbol,
                                                                 self._config)

    self._executors.passive = construct_mm_executor(self._parent._exchanges.passive,
                                                    execution_params_passive,
                                                    self._parent.ogs.passive,
                                                    self._logger)

    # Register fill callback
    self._parent.ogs.passive.register_fill_callback(self.on_passive_fill)

  def on_base_init(self, ioloop, queue, flow_sub):
    super().on_base_init(ioloop, queue, flow_sub)
    self.start_print_feature_status()

  def _has_pending_order(self):
    return any(og.has_pending_order(self._product) for og in self._parent.ogs)

  def on_position_update(self, timestamp):
    # self._on_book(timestamp, True)
    pass

  def _is_books_ready(self, limit_ns=2e9):
    # Check the timestamp and every book is ready
    for book in self._books:
      if book is None or not book.has_ask() or not book.has_bid():
        return False

    min_ts = min(book.timestamp for book in self._books)
    max_ts = max(book.timestamp for book in self._books)
    if max_ts - min_ts > limit_ns:
      return False
    return True

  def _is_books_sane(self):
    # Check the sanity of every book
    for book in self._books:
      if (book.has_ask() and book.has_bid() and book.ask0().price <= book.bid0().price):
        self._logger.warn('Book %s is insane: ask0(%f) <= bid0(%f)',
                          book,
                          book.ask0().price,
                          book.bid0().price)
        return False
    return True

  def on_book(self, exchange_id, book):
    if exchange_id == PassAgg.PASSIVE:
      self._books.passive = book
    elif exchange_id == PassAgg.AGGRESSIVE:
      self._books.aggressive = book
    else:
      raise ValueError(exchange_id)

    if not self._is_books_ready() or not self._is_books_sane():
      return

    # balancer
    # TODO(inkyu): It sucks
    midp = (self._books.passive.bid0().price + self._books.passive.ask0().price
            + self._books.aggressive.bid0().price + self._books.aggressive.ask0().price) / 4
    self._parent._balance_printer._prices[self._arb_product.product.base] = midp

    # pricer
    self._pricer.update(book_aggressive=self._books.aggressive, book_passive=self._books.passive)

    # sizer
    arb_balance = self._parent._balance_printer.get_balance(self._arb_product,
                                                            self._reserve_balance)
    self._sizer.update(book.timestamp, arb_balance)

    time_diff = abs(self._books.passive.timestamp - self._books.aggressive.timestamp)
    if abs(int(time_diff / 1e9)) >= 2:
      # Trust only feed within 2 second.
      return False

    if not self._pricer.ready or not self._executors.passive or FLAGS.dry_run:
      return

    passive_prices = self._pricer.calculate_price_passive(self._edge, pull_edge=self._pull_edge)
    passive_sizes = self._sizer.calculate_size_passive()
    if passive_sizes is None:
      self._logger.warn('[PASS:%s] Sizer is not ready.' % self._arb_product.product)
      return

    pass_fired = self._executors.passive.manage_mm_orders(sell_price=passive_prices.askp,
                                                          buy_price=passive_prices.bidp,
                                                          sell_qty=passive_sizes.sell_qty,
                                                          buy_qty=passive_sizes.buy_qty,
                                                          sell_price_pull=passive_prices.askp_pull,
                                                          buy_price_pull=passive_prices.bidp_pull)

    if pass_fired:
      time_diff = abs(self._books.passive.timestamp - self._books.aggressive.timestamp)
      # If there is no new passive order fired, do not cancel aggression orders.
      self.cancel_all_aggressive_orders()

    # Fire pending aggression orders
    expired_order_sides = []

    for order_side, pending_aggression in self._pending_aggressions.items():
      if pending_aggression is not None:
        fired = self._try_aggression(order_side=order_side,
                                     fill_price=pending_aggression.fill_price,
                                     fill_qty=pending_aggression.fill_qty)
        if fired or book.timestamp - pending_aggression.timestamp > 5e9:  # TODO(inkyu): Config
          expired_order_sides.append(order_side)
    for order_side in expired_order_sides:
      self._pending_aggressions[order_side] = None

  @property
  def _timestamp_last_book(self):
    timestamp = 0
    for book in self._books:
      if book and book.timestamp > timestamp:
        timestamp = self._books.passive.timestamp
    assert timestamp

    return timestamp

  def on_passive_fill(self, fill_price, fill_qty, order):
    # TODO(inkyu): Receive timestamp from the og
    timestamp = self._timestamp_last_book

    if order is None or order.product != self._product_passive or FLAGS.dry_run:
      return

    if not self._parent.ogs.aggressive:
      self._logger.error('OG is not ready. Skip aggression')
      return

    # sizer
    arb_balance = self._parent._balance_printer.get_balance(self._arb_product,
                                                            self._reserve_balance)
    self._sizer.update(timestamp, arb_balance)
    if order.order_side == OrderSide.BUY:
      self._sizer.on_passive_fill_buy(qty_buy=fill_qty, timestamp=timestamp)
    elif order.order_side == OrderSide.SELL:
      self._sizer.on_passive_fill_sell(qty_sell=fill_qty, timestamp=timestamp)
    else:
      raise ValueError(order.order_side)

    self._logger.info(
        '[PASS:%s] FILL %s %f@%f imbalance: %f' %
        (order.product, order.order_side, fill_qty, fill_price, self._sizer.get_imbalance() or 0.))

    self._logger.info('[AGG] bid0: %.8f ask0: %.8f' %
                      (self._books.aggressive.bid0().price, self._books.aggressive.ask0().price))

    order_side_agg = agg_order_side(order)

    self._try_aggression(order_side_agg, fill_price, fill_qty)

  def _try_aggression(self, order_side, fill_price, fill_qty):
    if order_side == OrderSide.BUY:
      return self.try_aggression_buy(fill_price, fill_qty)
    elif order_side == OrderSide.SELL:
      return self.try_aggression_sell(fill_price, fill_qty)
    else:
      raise ValueError('Invalid side: %s' % str(order_side))

  def on_agg_fill(self):
    pass

  def cancel_all_aggressive_orders(self):
    if FLAGS.dry_run:
      return False
    return self._parent.ogs.aggressive.cancel_product(self._arb_product.native_aggressive,
                                                      ignore_error=True)

  def prune_fire_qty(self, original_qty):
    prune_qty = original_qty
    if prune_qty < self._lot_size_min_agg:
      return False, original_qty, 0

    prune_qty = min(prune_qty, self._lot_size_max_agg)
    prune_qty = self._product_holders.aggressive.get_formatter().floor_qty(prune_qty)

    return True, original_qty, prune_qty

  def prune_fire_price(self, original_price, order_side):
    if order_side == OrderSide.BUY:
      return self._product_holders.aggressive.get_tick().floor(original_price)
    elif order_side == OrderSide.SELL:
      return self._product_holders.aggressive.get_tick().ceil(original_price)
    else:
      raise ValueError(order_side)

  def try_aggression_buy(self, fill_price, fill_qty):
    if FLAGS.dry_run:
      return False

    prices = self._pricer.calculate_price_aggressive_buy(fill_price, self._agg_edge)
    sizes = self._sizer.calculate_size_aggressive()
    if sizes is None:
      self._logger.warn('[AGG:%s] Sizer is not ready' % self._arb_product.product.symbol)
      return False

    do_buy, original_qty, buy_qty = self.prune_fire_qty(sizes.buy_qty)
    if not do_buy:
      self._logger.info('[AGG:%s] NO BUY: qty=%.8f < lot_size_min=%.8f',
                        self._arb_product.product.symbol,
                        original_qty,
                        self._lot_size_min_agg)
      return False
    buy_price = self.prune_fire_price(prices.posting_price, OrderSide.BUY)

    self._parent.ogs.aggressive.submit(self._product_aggressive,
                                       buy_price,
                                       buy_qty,
                                       OrderSide.BUY,
                                       OrderType.LIMIT)
    self._last_timestamp_agg_buy = self._timestamp_last_book
    self._logger.info('[AGG:%s] BUY  %.8f @ %.8f '
                      'fill=%.8f, askt_agg=%.8f, edge=%.2fbp',
                      self._arb_product.product.symbol,
                      buy_qty,
                      buy_price,
                      fill_price,
                      prices.true_book_price,
                      prices.edge * 10000.0)

    self._sizer.on_aggressive_fire_buy(qty_buy=buy_qty, timestamp=self._timestamp_last_book)
    return True

  def try_aggression_sell(self, fill_price, fill_qty):
    if FLAGS.dry_run:
      return False

    prices = self._pricer.calculate_price_aggressive_sell(fill_price, self._agg_edge)
    sizes = self._sizer.calculate_size_aggressive()
    if sizes is None:
      self._logger.warn('[AGG:%s] Sizer is not ready' % self._arb_product.product.symbol)
      return False

    do_sell, original_qty, sell_qty = self.prune_fire_qty(sizes.sell_qty)
    if not do_sell:
      self._logger.info('[AGG:%s] NO SELL: qty=%.8f < lot_size_min=%.8f',
                        self._arb_product.product.symbol,
                        original_qty,
                        self._lot_size_min_agg)
      return False
    sell_price = self.prune_fire_price(prices.posting_price, OrderSide.SELL)

    self._parent.ogs.aggressive.submit(self._product_aggressive,
                                       sell_price,
                                       sell_qty,
                                       OrderSide.SELL,
                                       OrderType.LIMIT)
    self._last_timestamp_agg_sell = self._timestamp_last_book
    self._logger.info('[AGG:%s] SELL %.8f @ %.8f '
                      'fill=%.8f, bidt_agg=%.8f, edge=%.2fbp',
                      self._arb_product.product.symbol,
                      sell_qty,
                      sell_price,
                      fill_price,
                      prices.true_book_price,
                      prices.edge * 10000.0)

    self._sizer.on_aggressive_fire_sell(qty_sell=sell_qty, timestamp=self._timestamp_last_book)
    return True


class SoftArbStrategy(object):
  def __init__(self, strategy_params):
    self._logger = logging.getLogger(__name__)

    self._feed = MultiFeedSubsystem()
    self._config = strategy_params['config']

    self._og_configs = strategy_params['og_configs']
    self._subscription_requests = PassAggExchangeMap._make(
        SubscriptionRequest.from_str(subscription_request_str)
        for subscription_request_str in strategy_params['subscription_requests'])

    # TODO(someone): Need some unified metaclass
    self._book_builder_clses = PassAggExchangeMap._make(
        get_flow_book_builder_cls(subscription_request)
        for subscription_request in self._subscription_requests)
    self._live_env_clses = PassAggExchangeMap._make(
        get_live_env_cls(subscription_request)
        for subscription_request in self._subscription_requests)
    self._exchanges = PassAggExchangeMap._make(
        subscription_request.exchange for subscription_request in self._subscription_requests)
    self._live_envs = PassAggExchangeMap(passive=None, aggressive=None)

    self._subreq_id_map = {
        self._subscription_requests[0]: PassAgg.PASSIVE,
        self._subscription_requests[1]: PassAgg.AGGRESSIVE,
    }

    self.ogs = PassAggExchangeMap(passive=None, aggressive=None)

    self._single_product_sub_strategies = {}

    self._arb_products = None
    self._balance_printer = None
    self._order_canceled_after_og_ready = False
    self._exit_attempt = 0
    self._exit_executor = None
    self._og_warn = [True, True]

  def _on_bnb_book(self, book):
    midp = (book.bid0().price + book.ask0().price) / 2
    self._balance_printer._prices[SharedCurrency.FromStrCurrency('BNB')] = midp

  def subscribe_feeds(self):
    common_config = self._config['common']
    self._arb_products = PassAggArbProducts(common_config['quote_currency'],
                                            self._config['products'].keys(),
                                            self._book_builder_clses[0].ProductType,
                                            self._book_builder_clses[1].ProductType)

    self._balance_printer = ArbInfoPnlPrinter(
        quote_currency=self._arb_products.shared.quote_currency,
        base_currencies=self._arb_products.shared.base_currencies)

    # TODO(inkyu): Is there a better way?
    if 'bnb_reserve' in common_config:
      bnb_reserve = float(common_config['bnb_reserve'])
      self._balance_printer._reserves[SharedCurrency.FromStrCurrency('BNB')] = bnb_reserve

    for symbol, config in self._config['products'].items():
      config = {**common_config, **config}
      arb_product = self._arb_products.by_symbol(symbol)
      self._single_product_sub_strategies[arb_product.product] = SingleProductSubStrategy(
          arb_product, self, config, logger=self._logger)

    self._feed.subscribe_book(self._arb_products.natives.passive.products,
                              functools.partial(self._on_book, PassAgg.PASSIVE))
    self._feed.subscribe_book(self._arb_products.natives.aggressive.products,
                              functools.partial(self._on_book, PassAgg.AGGRESSIVE))
    # TODO(inkyu): Clean up
    if 'binance' in self._subreq_id_map:
      product_bnb = BinanceProduct(base=BinanceCurrency.FromStrCurrency('BNB'),
                                   quote=BinanceCurrency.FromCurrency(
                                       self._arb_products.shared.quote_currency))
      self._feed.subscribe_book(product_bnb, self._on_bnb_book)

  def subscribe_live_envs(self, ioloop, queue, flow_sub):
    self._live_envs.passive = self._live_env_clses.passive(
        self._arb_products.natives.passive.currencies,
        self._arb_products.natives.passive.products,
        self._og_configs.passive,
        ioloop=ioloop,
        queue=queue,
        flow_sub=flow_sub)
    self._live_envs.passive.on_order_gateway_reset_callback.append(self.register_order_gateway)
    self._live_envs.aggressive = self._live_env_clses.aggressive(
        self._arb_products.natives.aggressive.currencies,
        self._arb_products.natives.aggressive.products,
        self._og_configs.aggressive,
        ioloop=ioloop,
        queue=queue,
        flow_sub=flow_sub)
    self._live_envs.aggressive.on_order_gateway_reset_callback.append(self.register_order_gateway)

  def register_order_gateway(self, subreq, order_gateway):
    self._logger.info('OG initialized: %s', subreq)
    subreq_id = self._subreq_id_map[subreq]
    if subreq_id == PassAgg.PASSIVE:
      self.ogs.passive = order_gateway
    elif subreq_id == PassAgg.AGGRESSIVE:
      self.ogs.aggressive = order_gateway
    if all(self.ogs):
      # cancel working orders first after OGs are ready.
      if not self._order_canceled_after_og_ready:
        self.ogs.passive.cancel_all()
        self.ogs.aggressive.cancel_all()
        self._logger.info('Orders canceled after OGs are ready.')
        self._order_canceled_after_og_ready = True
      for sub_strategy in self._single_product_sub_strategies.values():
        sub_strategy.on_og_reset()

    if not self._balance_printer.ready and all(self.ogs):
      self._balance_printer.set_ogs(self.ogs)
      self._balance_printer.print_balance()

    self._exit_executor = SimpleCancelAllExecutor()

  def start(self, ioloop, queue, flow_sub):
    self._ioloop = ioloop
    self._feed.start(ioloop, queue, flow_sub)
    for live_env in self._live_envs:
      live_env.start()
    self.print_balance(periodic=True)

  def _on_book(self, subreq_id, book):
    if not self.ogs.passive:
      if self._og_warn[0]:
        self._logger.info('%s OG not ready' % self._exchanges.passive)
        self._og_warn[0] = False
      return
    if not self.ogs.aggressive:
      if self._og_warn[1]:
        self._logger.info('%s OG not ready' % self._exchanges.passive)
        self._og_warn[1] = False
      return
    if self._exit_attempt > 0:
      return

    product = SharedProduct.FromProduct(book.product)
    self._single_product_sub_strategies[product].on_book(subreq_id, book)

  def print_balance(self, periodic=False):
    if self._balance_printer.ready:
      self._balance_printer.print_balance()
      print_working_orders(self._logger, self.ogs.passive)
      print_working_orders(self._logger, self.ogs.aggressive)
    if periodic:
      self._ioloop.add_timeout(datetime.timedelta(seconds=10), self.print_balance, True)

  def signal_handler(self, sig, frame):
    self._logger.warning('Caught signal: %s', sig)
    self._ioloop.add_callback_from_signal(self.on_exit)

  def on_exit(self):
    self._exit_attempt += 1
    if self._exit_attempt > 5:
      self._ioloop.add_callback(self._ioloop.stop)

    finished = True
    if self._exit_executor is not None:
      ogs = [self.ogs.passive, self.ogs.aggressive]
      self._logger.info('Exit callback executed.')
      finished = self._exit_executor.cancel_all_on_exit(ogs)

    if finished:
      self._ioloop.add_callback(self._ioloop.stop)
    else:
      timeout = datetime.timedelta(seconds=2)
      self._ioloop.add_timeout(timeout, self.on_exit)


def get_param_set(subreq_passive,
                  subreq_aggressive,
                  og_config_passive,
                  og_config_aggressive,
                  config_filename):
  config = load_clean_json(config_filename)
  return {
      'subscription_requests':
          PassAggExchangeMap(passive=subreq_passive, aggressive=subreq_aggressive),
      'og_configs':
          PassAggExchangeMap(passive=og_config_passive, aggressive=og_config_aggressive),
      'config':
          config,
  }


def main(argv):
  strategy_params = get_param_set(subreq_passive=FLAGS.subreq_passive,
                                  subreq_aggressive=FLAGS.subreq_aggressive,
                                  og_config_passive=FLAGS.og_config_passive,
                                  og_config_aggressive=FLAGS.og_config_aggressive,
                                  config_filename=FLAGS.config_filename)

  queue = SimpleQueue('pass_agg')
  flow_sub = flow_subscriber.from_queue(queue)
  ioloop = IOLoop.current()
  strategy = SoftArbStrategy(strategy_params)
  strategy.subscribe_feeds()
  strategy.subscribe_live_envs(ioloop, queue, flow_sub)  # What the..
  strategy.start(ioloop, queue, flow_sub)
  exit_after = datetime.timedelta(minutes=FLAGS.exit_after_min)
  ioloop.add_timeout(exit_after, strategy.on_exit)

  signal.signal(signal.SIGTERM, strategy.signal_handler)
  signal.signal(signal.SIGINT, strategy.signal_handler)

  ioloop.start()


if __name__ == '__main__':
  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
